package main

import (
	"fmt"
	"gitee.com/lianyi_work/redis_ope"
	"github.com/miekg/dns"
	cmap "github.com/orcaman/concurrent-map/v2"
	log "github.com/sirupsen/logrus"
	"github.com/spf13/viper"
	"root_dns_server/logs"
)

type Nameserver struct {
	NS string `yaml:"ns"`
	A  string `yaml:"a"`
}

type DomainConfig struct {
	Domain      string       `yaml:"domain"`
	Nameservers []Nameserver `yaml:"nameservers"`
}

type Config struct {
	RedisAddr string         `yaml:"redis_addr"`
	RedisAuth string         `yaml:"redis_auth"`
	Config    []DomainConfig `yaml:"config"`
}

type ZoneConfig struct {
	domain string
	rr     dns.RR
}

func NewRR(s string) dns.RR { r, _ := dns.NewRR(s); return r }

var (
	zones_ns = cmap.New[[]*ZoneConfig]()
	zones_a  = cmap.New[[]dns.RR]()
)

func init() {
	viper.SetConfigName("config")
	viper.SetConfigType("yaml")
	viper.AddConfigPath(".")
	if err := viper.ReadInConfig(); err != nil {
		log.Fatalf("Error reading config file, %s", err)
	}
}

func main() {
	reloadZoneConfig()

	go func() {
		srv := &dns.Server{Addr: ":53", Net: "udp"}
		if err := srv.ListenAndServe(); err != nil {
			log.Fatalf("Failed to set udp listener %s\n", err.Error())
		}
	}()

	RunRemoteConfig()
	go runStatisticTask()

	select {}
}

func reloadZoneConfig() error {
	var config Config
	if err := viper.Unmarshal(&config); err != nil {
		log.Fatalf("Unable to unmarshal config, %v", err)
	}
	for _, domainConfig := range config.Config {
		fmt.Printf("Domain: %s\n", domainConfig.Domain)
		ns_value := make([]*ZoneConfig, 0)
		a_value := make([]dns.RR, 0)
		for _, ns := range domainConfig.Nameservers {
			fmt.Printf("  NS: %s, A: %s\n", ns.NS, ns.A)
			ns_value = append(ns_value, &ZoneConfig{domain: domainConfig.Domain, rr: NewRR(fmt.Sprintf("%s 600 IN NS %s", domainConfig.Domain, ns.NS))})
			a_value = append(a_value, NewRR(fmt.Sprintf("%s 600 IN A %s", ns.NS, ns.A)))
		}
		zones_ns.Set(domainConfig.Domain, ns_value)
		zones_a.Set(domainConfig.Domain, a_value)
	}

	dns.DefaultServeMux = dns.NewServeMux()
	for z, zone_configs := range zones_ns.Items() {
		dns.HandleFunc(z, func(w dns.ResponseWriter, r *dns.Msg) {
			m := new(dns.Msg)
			m.SetReply(r)
			m.Authoritative = true

			domain_name := ""
			for _, zone_ns := range zone_configs {
				domain_name = zone_ns.domain
				rr := []dns.RR{zone_ns.rr}
				m.Ns = append(m.Ns, rr...)
			}

			if a_rr, exist := zones_a.Get(domain_name); exist {
				m.Extra = append(m.Extra, a_rr...)
			}

			//add root log
			logs.AddAllLogData(r.Question[0].String(), w.RemoteAddr().String(), r.Question[0].Name, dns.TypeToString[r.Question[0].Qtype])

			w.WriteMsg(m)
		})
	}

	return nil
}

func runStatisticTask() {
	//redis_ope.Connect()
	node_ip := "root"
	es_addr, err := redis_ope.GetValue("dns-rec:sys-config:elk:addr")
	if err != nil {
		log.Errorf("get elasticsearch_addr from redis failed: %s", err)
	}
	es_pwd, err := redis_ope.GetValue("dns-rec:sys-config:elk:password")
	if err != nil {
		log.Errorf("get elasticsearch_pwd from redis failed: %s", err)
	}
	log.Tracef("start statistic server...")
	log.Tracef("statistic server info: %s, %s", es_addr, es_pwd)
	go logs.Run(node_ip, "http://"+es_addr, es_pwd, "")
}
